from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory

from a0base.base_llm import pop_llm

# 0. 全局内存存储
store = {}


def get_session_history(session_id: str) -> InMemoryChatMessageHistory:
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
    return store[session_id]


# 1.  chain 定义
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是乐于助人的助手。"),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{input}")
])
chain = prompt | pop_llm

# 2. 包装
chain_with_history = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="history"
)

# 3. 多轮调用
config = {"configurable": {"session_id": "demo"}}
print(chain_with_history.invoke({"input": "我叫小明"}, config))
print(chain_with_history.invoke({"input": "我叫什么？"}, config))
